Greptile Summary

This PR enables 2D NVFP4 quantization in columnwise-only mode, needed by HybridQuantizer for MXFP8-fwd + NVFP4-bwd on weights. Two separate code paths are updated: the optimized TMA-based kernel (quantize_transpose_nvfp4_2D_kernel) gains a RETURN_ROWWISE template boolean that gates out the rowwise scaling pass and data store at compile time, and the fallback blockwise kernel (block_scaled_1d_cast_transpose_kernel) gains a new "Step 2.5" that re-runs the load + local-amax + 2D warp/smem reduction from Step 2 when only the transposed output is requested, populating amax_smem for Step 3 without performing the rowwise scale/quantize/store writes.

• Dispatch (quantize.cuh): use_optimized_kernel is extended to route BF16 + 32-aligned + columnwise-only + 2D requests to the TMA kernel instead of falling back to the blockwise path.
• Optimized path (quantize_transpose_nvfp4.cuh): RETURN_ROWWISE template parameter added; scale_stride, scales_ptr, and tensor_map_output are all guarded so the kernel safely no-ops on their rowwise uses when RETURN_ROWWISE=false.
• Fallback path (quantize_transpose_vector_blockwise_fp4.cu): New Step 2.5 correctly mirrors the 2D amax-reduction pass from Step 2, with matching __syncthreads() barriers, so amax_smem is fully populated for Step 3 in the columnwise-only 2D case.

Confidence Score: 5/5

The changes are well-scoped: new code paths are gated by constexpr template parameters or function-level preconditions, and the new Step 2.5 amax-only pass faithfully mirrors the existing 2D reduction without touching any existing output writes.

All access sites for rowwise pointers and tensor maps in the columnwise-only case are guarded at compile time (constexpr if). Synchronization in Step 2.5 is correct: amax_smem_red is written and synced before reduction, and amax_smem is synced before Step 3 reads it. The bitwise-equality test covers both the TMA kernel (BF16 + aligned) and the blockwise fallback (float32 + non-aligned), leaving minimal room for undetected regressions.

No files require special attention; the most complex change (Step 2.5 in the .cu file) is validated by the new test.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["quantize_fwd/bwd_helper\n(quantize.cuh)"] --> B{use_optimized_kernel?}
    B -->|"BF16 + rows%32==0 + cols%32==0\n+ (has_data OR (has_colwise_data AND 2D))"| C["quantize_transpose\n(quantize_transpose_nvfp4.cuh)"]
    B -->|else| D["quantize_transpose_vector_blockwise_fp4\n(fallback kernel)"]
    C --> E{use_2d_quantization?}
    E -->|true| F["quantize_transpose_nvfp4_2D_kernel\nRETURN_ROWWISE=bool\nRETURN_TRANSPOSE=bool"]
    E -->|false| G["quantize_transpose_nvfp4_kernel (1D)\nguarded: return_rowwise||use_2d must be true"]
    F --> H{RETURN_ROWWISE?}
    H -->|true| I["Step: 2D amax + rowwise scale/quantize/store\n+ TMA cp_async rowwise data"]
    H -->|false| J["Step: 2D amax only\n(block_amax_matrix populated)"]
    I --> K{RETURN_TRANSPOSE?}
    J --> K
    K -->|true| L["Colwise scale/quantize/store\n+ TMA cp_async colwise data"]
    D --> M{kReturnIdentity\n&& kIs2DBlockScaling?}
    M -->|"kReturnIdentity=true"| N["Step 2: full rowwise pass\n(populates amax_smem)"]
    M -->|"!kReturnIdentity\n&& kIs2DBlockScaling=true"| O["Step 2.5 (NEW): amax-only pass\nload smem → amax_smem_red → amax_smem"]
    N --> P["Step 3: transpose cast+store\n(reads amax_smem for 2D)"]
    O --> P

_{Reviews (2): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

This is just the fallback kernel being changed. Does the main kernel already support this?

This is just the fallback kernel being changed. Does the main kernel already support this?

Thanks for the catch. The main kernel does not support if as well. Enabled in f7953dd